#!/usr/bin/env python
# encoding: utf-8
"""
convertGenepop.py

Created by Brant Faircloth on 2007-10-30.
Copyright (c) 2007 Brant C. Faircloth. All rights reserved.

Convert cervus (csv-formatted) data to genepop format (3 number/fragment sizes)

"""

import csv, numpy, time, pdb

#infile = ''
infile  = ''
outfile = ''
# columns to drop from the population
drop = [1,2]
# columns containin pop info
pop = ''
# header info to remove
headD = [0,3]

def remove(arr, index, dim=0):  
    """
    from "Perry" on nummarray listserv
    
    Remove subarray at index and dimension specified
    
    Usage:
    to remove the 6th row:
    x = remove(x, 5)
    to remove the 5th column:
    x = remove(x, 4, dim=1)  
    """
    #pdb.set_trace()
    # make selected dimension the first one
    swparr = numpy.swapaxes(arr, 0, dim)
    indices = range(swparr.shape[0])
    if type(index) == int:
        del indices[index]
    elif type(index) == list:
        # this will only hold 2 numbers to form 1:2 or whatever
        del indices[index[0]:index[1]+1]
    newarray = swparr[indices]
    # reorder axes as they originally appeared.
    return numpy.swapaxes(newarray, 0, dim)

class excelSingleSpace:
    """class for csv module to work correctly"""
    delimiter = ','
    quotechar = '"'
    escapechar = None
    doublequote = True
    skipinitialspace = False
    lineterminator = '\r'


def read():
    '''read data into a large numpy array (l by w of data)'''
    handle  = open(infile,'rU')
    c       = csv.reader(handle, delimiter=',', dialect = excelSingleSpace)
    indiv   = int(c.next()[0])
    loci     = int(c.next()[0])
    # skip remaining header
    header = c.next()
    # create an array to hold the data using the first row of real data cast to int
    #geno = [int(i) for i in c.next()[3:]]
    geno = [int(i) for i in c.next()]
    geno = numpy.array(geno)
    for row in c:
        # cast list of strings to ints
        row = [int(i) for i in row]
        # convert list to array (slicing out other info)
        vector = numpy.array(row)
        # vertically stack with original array
        geno = numpy.vstack((geno,vector))
    #pdb.set_trace()
    return geno,header

def head(header):
    del header[headD[0]:headD[1]+1]
    # keep only odd members of header
    temp = []
    [temp.append(header[i]) for i in range(0,len(header),2)]
    return temp

def convert(geno):
    geno = remove(geno,drop,dim=1)
    return geno

def write(header,geno):
    file = open(outfile,'w')
    dt = time.strftime('%m/%d/%Y %H:%M:%S')
    file.write(('Genepop data file converted to genepop 3-integer format on %s\n') % (dt))
    for locus in header:
        file.write(('%s\n') % (locus))
    file.write('pop\n')
    #pdb.set_trace()
    for row in geno:
        file.write(('%s ,') % (row[0]))
        for allele in range(1,len(row),2):
            if row[allele] in [-9,'*',0,'-9','0'] or row[allele+1] in [-9,'*',0,'-9','0']:
                file.write(' 000000')
            else:
                file.write((' %s%s') % (row[allele],row[allele+1]))
        file.write('\n')
    file.close()
    
def main():
    geno,header     = read()
    # remove info from header
    header          = head(header)
    geno            = convert(geno)
    write(header,geno)


if __name__ == '__main__':
    main()

